import torch
import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
from typing import Dict, List, Optional
import seaborn as sns


class FeatureVisualizer:
    """特征可视化器"""
    
    def __init__(self, save_dir: Path):
        self.save_dir = Path(save_dir)
        self.save_dir.mkdir(parents=True, exist_ok=True)
        self.hooks = []
        self.feature_maps = {}
        
    def register_hooks(self, model: torch.nn.Module, target_layers: List[str]):
        """注册前向钩子来捕获特征图"""
        
        def get_activation(name):
            def hook(model, input, output):
                if isinstance(output, torch.Tensor):
                    self.feature_maps[name] = output.detach().cpu()
                    print(f"捕获特征图: {name}, 形状: {output.shape}")
                elif isinstance(output, (tuple, list)) and len(output) > 0:
                    if isinstance(output[0], torch.Tensor):
                        self.feature_maps[name] = output[0].detach().cpu()
                        print(f"捕获特征图 (第一个): {name}, 形状: {output[0].shape}")
            return hook
        
        registered_count = 0
        for name, module in model.named_modules():
            if name in target_layers:
                try:
                    hook = module.register_forward_hook(get_activation(name))
                    self.hooks.append(hook)
                    registered_count += 1
                    print(f"注册特征钩子: {name}")
                except Exception as e:
                    print(f"注册特征钩子失败 {name}: {e}")
        
        print(f"总共注册了 {registered_count} 个特征钩子")

    def remove_hooks(self):
        """移除所有钩子"""
        for hook in self.hooks:
            hook.remove()
        self.hooks.clear()
        
    def visualize_feature_maps(self, layer_name: str, max_channels: int = 16):
        """可视化指定层的特征图"""
        if layer_name not in self.feature_maps:
            print(f"警告: 层 {layer_name} 没有特征图数据")
            return
            
        feature_map = self.feature_maps[layer_name]
        # 检查特征图维度
        if len(feature_map.shape) < 4:
            print(f"警告: 层 {layer_name} 特征图维度不足: {feature_map.shape}")
            return
        
        batch_size, channels, height, width = feature_map.shape
        print(f"可视化层 {layer_name}: {channels} 通道, {height}x{width}")
                
        # 选择要可视化的通道数
        num_channels = min(channels, max_channels)
        
        # 创建子图
        cols = 4
        rows = (num_channels + cols - 1) // cols
        fig, axes = plt.subplots(rows, cols, figsize=(15, 4 * rows))
        fig.suptitle(f'Feature Maps - {layer_name}', fontsize=16)
        
        if rows == 1:
            axes = axes.reshape(1, -1)
        elif rows == 0:
            return
        
        for i in range(num_channels):
            row = i // cols
            col = i % cols
            
            # 获取特征图 (取第一个batch)
            feature = feature_map[0, i].numpy()

            # 检查特征是否为空或全为零
            if feature.size == 0:
                axes[row, col].text(0.5, 0.5, 'Empty', ha='center', va='center')
                axes[row, col].set_title(f'Channel {i}')
                axes[row, col].axis('off')
                continue

            # 可视化
            try:
                im = axes[row, col].imshow(feature, cmap='viridis')
                axes[row, col].set_title(f'Ch {i} ({feature.min():.3f}, {feature.max():.3f})')
                axes[row, col].axis('off')
                plt.colorbar(im, ax=axes[row, col], fraction=0.046, pad=0.04)
            except Exception as e:
                axes[row, col].text(0.5, 0.5, f'Error: {str(e)[:20]}', ha='center', va='center')
                axes[row, col].set_title(f'Channel {i}')
                axes[row, col].axis('off')
        
        # 隐藏多余的子图
        for i in range(num_channels, rows * cols):
            row = i // cols
            col = i % cols
            axes[row, col].axis('off')
        
        plt.tight_layout()
        save_path = self.save_dir / f'feature_maps_{layer_name.replace(".", "_")}.png'
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.close()
        
        print(f"特征图已保存: {save_path}")
        
    def visualize_feature_statistics(self):
        """可视化特征统计信息"""
        if not self.feature_maps:
            print("没有特征图数据可用于统计")
            return
            
        stats = {}
        for layer_name, feature_map in self.feature_maps.items():
            feature_flat = feature_map.flatten().numpy()
            # 检查数据有效性
            if len(feature_flat) == 0:
                print(f"警告: 层 {layer_name} 特征图为空")
                continue

            
            # 计算稀疏性 - 使用更合理的阈值，稀疏性 = 接近零的激活值数量 / 总激活值数量
            zero_threshold = 1e-6  # 接近零的阈值
            sparsity_ratio = np.mean(np.abs(feature_flat) < zero_threshold)
            
            stats[layer_name] = {
                'mean': float(np.mean(feature_flat)),
                'std': float(np.std(feature_flat)),
                'min': float(np.min(feature_flat)),
                'max': float(np.max(feature_flat)),
                'sparsity': float(sparsity_ratio),
                'non_zero_ratio': float(1.0 - sparsity_ratio),
                'median': float(np.median(feature_flat)),
                'l1_norm': float(np.mean(np.abs(feature_flat))),
                'shape': list(feature_map.shape)
            }
            
            print(f"层 {layer_name}: 均值={stats[layer_name]['mean']:.4f}, "
                    f"稀疏性={stats[layer_name]['sparsity']:.4f}")
            
        if not stats:
            print("没有有效的特征统计数据")
            return {}
        
        # 绘制统计图
        try:
            self._plot_feature_statistics(stats)
        except Exception as e:
            print(f"绘制特征统计图失败: {e}")
        
        return stats
    
    def _plot_feature_statistics(self, stats: Dict):
        """绘制特征统计图"""
        fig, axes = plt.subplots(2, 3, figsize=(18, 10))
        fig.suptitle('Feature Map Statistics', fontsize=16)
        
        layer_names = list(stats.keys())
        short_names = [name.split('.')[-1][:10] for name in layer_names]  # 使用短名称
        
        metrics = ['mean', 'std', 'sparsity', 'non_zero_ratio', 'l1_norm', 'median']
        titles = ['Mean Activation', 'Standard Deviation', 'Sparsity Ratio', 
                 'Non-zero Ratio', 'L1 Norm', 'Median']
        
        for idx, (metric, title) in enumerate(zip(metrics, titles)):
            row = idx // 3
            col = idx % 3
            
            values = [stats[name].get(metric, 0) for name in layer_names]
            
            # 检查数据有效性
            if all(v == 0 for v in values):
                axes[row, col].text(0.5, 0.5, f'No {metric} data', 
                                  ha='center', va='center', transform=axes[row, col].transAxes)
            else:
                bars = axes[row, col].bar(range(len(layer_names)), values)
                
                # 为每个条形添加数值标签
                for i, (bar, val) in enumerate(zip(bars, values)):
                    axes[row, col].text(bar.get_x() + bar.get_width()/2, bar.get_height(),
                                      f'{val:.3f}', ha='center', va='bottom', fontsize=8)
            
            axes[row, col].set_title(title)
            axes[row, col].set_xticks(range(len(layer_names)))
            axes[row, col].set_xticklabels(short_names, rotation=45, ha='right')
            axes[row, col].tick_params(axis='x', labelsize=8)
            axes[row, col].grid(True, alpha=0.3)
        
        plt.tight_layout()
        save_path = self.save_dir / 'feature_statistics.png'
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.close()
        
        print(f"特征统计图已保存: {save_path}")
        
        # 创建热图
        self._create_feature_heatmap(stats)

    def _create_feature_heatmap(self, stats: Dict):
        """创建特征统计热图"""
        try:
            layer_names = list(stats.keys())
            metrics = ['mean', 'std', 'sparsity', 'non_zero_ratio', 'l1_norm', 'median']
            
            # 创建数据矩阵
            data_matrix = []
            for layer_name in layer_names:
                row = [stats[layer_name].get(metric, 0) for metric in metrics]
                data_matrix.append(row)
            
            data_matrix = np.array(data_matrix)
            
            if data_matrix.size == 0:
                print("没有数据用于创建特征热图")
                return
            
            # 标准化数据
            std_vals = data_matrix.std(axis=0)
            std_vals[std_vals == 0] = 1
            normalized_data = (data_matrix - data_matrix.mean(axis=0)) / std_vals
            
            plt.figure(figsize=(10, max(8, len(layer_names) * 0.5)))
            
            short_names = [name.split('.')[-1][:15] for name in layer_names]
            
            sns.heatmap(normalized_data, 
                       xticklabels=metrics,
                       yticklabels=short_names,
                       cmap='RdBu_r', center=0, annot=True, fmt='.2f')
            
            plt.title('Feature Statistics Heatmap (Normalized)')
            plt.xlabel('Metrics')
            plt.ylabel('Layers')
            plt.tight_layout()
            
            save_path = self.save_dir / 'feature_statistics_heatmap.png'
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            plt.close()
            
            print(f"特征统计热图已保存: {save_path}")
        except Exception as e:
            print(f"创建特征热图失败: {e}")